package scu.maqiang.electromagnetics.homogenization;

import scu.maqiang.fes.BVPType;
import scu.maqiang.fes.FES2T31;
import scu.maqiang.mesh.*;
import scu.maqiang.numeric.Complex;
import scu.maqiang.numeric.IterSCSolver;
import scu.maqiang.numeric.SCMatrix;
import scu.maqiang.numeric.SRMatrix;

import java.util.HashSet;

public class ScatteringHomoSolution {
    public static void main(String[] args) {
        Mesh2T3 mesh = new Mesh2T3().square2D(192, 192, T3Type.Left);
        mesh.scale(1.5, 1.5);
        mesh.translate(-0.25, -0.25);
        mesh.toTecplot("FineMesh.dat");

        int N = 8;
        double eps = 1.0 / N;
        int DielDomainLabel = 20;
        ScalarFunc dielDomainFunc = (xy, label, param) -> Math.abs(xy[0] - 0.5) < 0.5 && Math.abs(xy[1] - 0.5) < 0.5?1.0:-1.0;
        mesh.setDomainLabel(dielDomainFunc, null, DielDomainLabel);

        int PMLDomainLabel = 10;
        ScalarFunc PMLDomainFunc = (xy, label, param) -> Math.abs(xy[0] - 0.5) < 0.625 && Math.abs(xy[1] - 0.5) < 0.625?-1.0:1.0;
        mesh.setDomainLabel(PMLDomainFunc, null, PMLDomainLabel);

        int outerBoundaryLabel = 0;
        ScalarFunc outerBoundaryFunc = (xy, label, param) -> 1.0;
        mesh.setBoundaryLabel(outerBoundaryFunc, null, outerBoundaryLabel);

        int interfaceBoundaryLabel = 6;
        ScalarFunc interBoundaryFunc = (xy, label, param) -> {
            if(Math.abs(Math.abs(xy[0] - 0.5) - 0.625) < 1.0e-6) {
                return Math.abs(xy[1] - 0.5) < 0.625?1.0:-1.0;
            }
            if(Math.abs(Math.abs(xy[1] - 0.5) - 0.625) < 1.0e-8) {
                return Math.abs(xy[0] - 0.5) < 0.625?1.0:-1.0;
            }
            return -1.0;
        };
        mesh.addInterfaceBoundary(interBoundaryFunc, null, interfaceBoundaryLabel);
        System.out.println(mesh.extractBoundaryNodes(interfaceBoundaryLabel).size());
        mesh.saveMesh(MeshFormat.msh, "HomoMesh.msh");


        double c0  = 3*1e8;           //m/sec, velocity of light in free space
        double nu0 = 120*Math.PI;          // ohm, intrinsic impedance of the free space
        double e0  = (1e-9)/(36*Math.PI);  // F/m, permittivity of free space
        double mu0 = 4*Math.PI*1e-7;       // H/m, permeability of free space

        // Input parameters
        //******************
        double freq = 300;          // MHz, frequency
        freq *= 1e6;     // Hz, frequency
        double lambda0 = c0/freq;      // meter, wavelength
        double k0  = 2*Math.PI/lambda0; // 1/meter, wavenumber
        double omg  = 2*Math.PI*freq;    // rad/sec, radial frequency

//        Map<Integer, Complex[]> pmlCoordinates = mesh.setPMLCoordinates(k0, new int[]{interfaceBoundaryLabel}, new int[]{outerBoundaryLabel}, PMLDomainLabel);
//        for(var entry: pmlCoordinates.entrySet()) {
//            System.out.print(entry.getKey() + "\t");
//            Complex[] xyz = entry.getValue();
//            for(int i = 0; i < xyz.length; i++) {
//                System.out.print(xyz[i] + "\t");
//            }
//            System.out.println();
//        }

        //入射角
        double phii = 0;            // degree, angle of incident field
        double varphi = phii * Math.PI/180;  // radian, angle of incident field
        //入射波函数的负值
        ScalarCFunc incidentWaveFunc = (xy, label, param) -> {
            Complex result = new Complex(0.0, 0.0);
            result.real =  -Math.cos(k0 * (xy[0] * Math.cos(varphi) + xy[1] * Math.sin(varphi)));
            result.image = -Math.sin(k0 * (xy[0] * Math.cos(varphi) + xy[1] * Math.sin(varphi)));
            return result;
        };

        double mur = 1; // relative permittivity of object (used if scat_type='diel')
        double er = 4;  // dielectric constant of dielectric object (used if scat_type='diel')


        FES2T31 fs = new FES2T31(mesh);
        SRMatrix A = new SRMatrix(fs.GetNdof());
        SCMatrix A_PML = new SCMatrix(fs.GetNdof());

        //组装介电区域左端矩阵
        fs.assembleStiff(new double[]{1.0/mur * 0.05892}, BVPType.COMMON, A, DielDomainLabel);
        fs.assembleMass(new double[]{-k0 * k0 * 1.75}, BVPType.CONSIST_MASS, A, DielDomainLabel);

        Complex[] uint = mesh.functionInMesh(incidentWaveFunc, null);
        System.out.println("OK!");
        //组装右端源项1. 源项只在电介质内有效
//        Complex[] RHS = Complex.zeros(fs.GetNdof());
//        A.mul(uint, RHS);

//        for(int i = 0; i < fs.GetNdof(); i++) {
//            System.out.println(i + "\t" + RHS[i]);
//        }

        //在矩阵A上增加自由空间区域整体矩阵
        fs.assembleStiff(new double[]{1.0/mur}, BVPType.COMMON, A, 0);
        fs.assembleMass(new double[]{-k0 * k0 * 1.0}, BVPType.CONSIST_MASS, A, 0);

        //组装右端项2. 如果这样组装的话，最终会得到和书中一样的结果, 但觉得这样组装有问题, 因为在边界上会组装到在自由空间中的一部分内容
        Complex[] RHS = Complex.zeros(fs.GetNdof());
        A.mul(uint, RHS);
        Complex[] newRHS = Complex.zeros(fs.GetNdof());
        HashSet<Integer> nodeSet = mesh.extractDomainNodes(DielDomainLabel);
        for(int idx:nodeSet) {
            newRHS[idx] = RHS[idx];
        }
        RHS = newRHS;

        System.out.println("OK!");
        //组装PML区域整体矩阵
        fs.assemblePMLStiffMass(mur, k0, -k0 * k0 * 1, BVPType.COMMON, BVPType.CONSIST_MASS,
                A_PML, interfaceBoundaryLabel, outerBoundaryLabel, PMLDomainLabel);

        //得到所有区域整体矩阵
        SCMatrix K = A_PML.add(1.0, A);
        System.out.println("OK!");
        //K.sort();
        //System.out.println(K.display());
        //Complex[] RHS = Complex.zeros(fs.GetNdof());

//        ScalarCFunc innerBdFunc = (xy, label, param) -> {
//            Complex result = new Complex(0.0, 0.0);
//            result.real = -Math.cos(k0 * xy[0]);
//            result.image = -Math.sin(k0 * xy[0]);
//            return result;
//        };
//        Complex[] uint = mesh.functionInMesh(innerBdFunc, null);
//        Complex[] RHS = K.mul(uint);



//        fs.applyBC_MBN(K, Direct.All, outerBoundaryLabel);
//        fs.applyBC_MBN(RHS, Direct.All, new Complex(0.0, 0.0), outerBoundaryLabel);
//
//        fs.applyBC_MBN(K, Direct.All, interfaceBoundaryLabel);
//        fs.applyBC_MBN(RHS, Direct.All, new ScalarCFunc[]{innerBdFunc}, interfaceBoundaryLabel);

        IterSCSolver solver = new IterSCSolver(K);
        Complex[] x = Complex.zeros(fs.GetNdof());
        solver.PCGSSSOR_2(RHS, x, 1.5, 2);

        double[][] result = new double[6][];
        result[0] = Complex.Real(x);
        result[1] = Complex.Image(x);
        result[2] = Complex.abs(x);
        Complex[] utot = Complex.add(x, -1.0, uint);
        result[3] = Complex.Real(utot);
        result[4] = Complex.Image(utot);
        result[5] = Complex.abs(utot);
        mesh.toTecplot("Scattering_compositeHomo2.dat", result);
    }
}
